"""
财务图表模块 - 提供财务指标图表生成功能
"""

import pandas as pd
import plotly.graph_objects as go
from typing import List, Dict, Any, Optional

def create_financial_chart(
    df: pd.DataFrame, 
    x_col: str, 
    y_cols: List[str], 
    title: str = "财务指标", 
    height: int = 500,
    bar_cols: Optional[List[str]] = None
) -> go.Figure:
    """
    创建财务指标图表
    
    参数:
        df: 包含财务数据的DataFrame
        x_col: X轴列名
        y_cols: Y轴列名列表
        title: 图表标题
        height: 图表高度
        bar_cols: 使用柱状图显示的列名列表，默认为None（全部使用线图）
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = [x_col] + y_cols
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 如果未指定柱状图列，则默认为空列表
    if bar_cols is None:
        bar_cols = []
    
    # 创建图表
    fig = go.Figure()
    
    # 添加每个指标的线或柱
    colors = ['blue', 'red', 'green', 'purple', 'orange', 'cyan', 'magenta']
    
    for i, col in enumerate(y_cols):
        color = colors[i % len(colors)]
        
        if col in bar_cols:
            # 使用柱状图
            fig.add_trace(go.Bar(
                x=df[x_col],
                y=df[col],
                name=col,
                marker_color=color
            ))
        else:
            # 使用线图
            fig.add_trace(go.Scatter(
                x=df[x_col],
                y=df[col],
                mode='lines+markers',
                name=col,
                line=dict(color=color)
            ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        yaxis_title="数值",
        height=height,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig

def create_income_statement_chart(
    df: pd.DataFrame, 
    title: str = "收入与利润趋势", 
    height: int = 500
) -> go.Figure:
    """
    创建利润表图表
    
    参数:
        df: 包含利润表数据的DataFrame，必须包含 '报告期', '营业收入', '净利润' 列
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['报告期', '营业收入', '净利润']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加营业收入柱状图
    fig.add_trace(go.Bar(
        x=df['报告期'],
        y=df['营业收入'],
        name='营业收入',
        marker_color='blue'
    ))
    
    # 添加净利润线图
    fig.add_trace(go.Scatter(
        x=df['报告期'],
        y=df['净利润'],
        mode='lines+markers',
        name='净利润',
        line=dict(color='red')
    ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="报告期",
        yaxis_title="金额",
        height=height,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig

def create_balance_sheet_chart(
    df: pd.DataFrame, 
    title: str = "资产负债结构", 
    height: int = 500
) -> go.Figure:
    """
    创建资产负债表图表
    
    参数:
        df: 包含资产负债表数据的DataFrame，必须包含 '报告期', '资产总计', '负债合计', '所有者权益合计' 列
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['报告期', '资产总计', '负债合计', '所有者权益合计']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加资产总计柱状图
    fig.add_trace(go.Bar(
        x=df['报告期'],
        y=df['资产总计'],
        name='资产总计',
        marker_color='blue'
    ))
    
    # 添加负债合计柱状图
    fig.add_trace(go.Bar(
        x=df['报告期'],
        y=df['负债合计'],
        name='负债合计',
        marker_color='red'
    ))
    
    # 添加所有者权益合计柱状图
    fig.add_trace(go.Bar(
        x=df['报告期'],
        y=df['所有者权益合计'],
        name='所有者权益合计',
        marker_color='green'
    ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="报告期",
        yaxis_title="金额",
        height=height,
        barmode='group',
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig

def create_cash_flow_chart(
    df: pd.DataFrame, 
    title: str = "现金流量趋势", 
    height: int = 500
) -> go.Figure:
    """
    创建现金流量表图表
    
    参数:
        df: 包含现金流量表数据的DataFrame，必须包含 '报告期' 列和至少一个现金流量列
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    if '报告期' not in df.columns:
        raise ValueError("数据缺少必要列: 报告期")
    
    # 查找现金流量相关列
    cash_flow_cols = [col for col in df.columns if '现金流' in col or '经营活动' in col or '投资活动' in col or '筹资活动' in col]
    
    if not cash_flow_cols:
        raise ValueError("数据中未找到现金流量相关列")
    
    # 创建图表
    fig = go.Figure()
    
    # 添加每个现金流量列的线图
    colors = ['blue', 'red', 'green', 'purple', 'orange']
    
    for i, col in enumerate(cash_flow_cols):
        color = colors[i % len(colors)]
        
        fig.add_trace(go.Scatter(
            x=df['报告期'],
            y=df[col],
            mode='lines+markers',
            name=col,
            line=dict(color=color)
        ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="报告期",
        yaxis_title="金额",
        height=height,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    return fig